
import tushare as ts
from pandas import DataFrame
import matplotlib.pyplot as plt

"""
获取回测周期内的年度财报数据，组合成一个dict数据结构，key是股票代码，
value是一个按照报告发布日期排序的列表，列表内的元素也是一个dict
{'eps': 每股收益, 'announced_date': 公告日期}
"""
def getCodeReports():

    # 这个tuple包含了三个元素，前两个分别是用来获取年度财报时的参数年份和季度，
    # 后面是一个发布年度财报的年份，因为一般情况下发布财报都是在第二年的4月底之前，
    # 所以这个年份比财报的年份晚一年
    reportDateTuples = [(2013, 4, 2014), (2014, 4, 2015)]

    # 要返回的数据结构
    '''
    {
        'code':[
            {
                'eps':0.56,
                'announced_date': '2018-09-26'
            }
        ]
    }
    '''
    codeReportDict = dict()

    # 循环获取所有指定报告期的数据
    for reportDateTuple in reportDateTuples:
        # 从Tushare获取年报数据
        dfReports = ts.get_report_data(reportDateTuple[0], reportDateTuple[1])
        print(dfReports)

        # 只需要股票代码、每股收益和公告日期三个字段
        codes = dfReports['code']
        epses = dfReports['eps']
        announcedDates = dfReports['report_date']

        # 这个是报告发布时的年度
        announcedYear = str(reportDateTuple[2])

        # 拿到已经缓存的股票代码集合
        codesOfCachedReports = set(codeReportDict.keys())

        # 循环获取所有数据
        for reportIndex in dfReports.index:
            code = codes[reportIndex]
            eps = epses[reportIndex]
            announcedDate = announcedDates[reportIndex]
            print('%s %5.2f %d %s' % (code, eps, reportDateTuple[0], announcedDate), flush=True)

            # 如果eps是非数字，或者发布日期的月份超过了4月，就不作处理，因为股票在上市前
            # 也会发布财报，那么这个财报的发布日期可能不是定期报告所规定的时间范围，
            # 那么对这种上市之前的数据暂时不予处理
            if str(eps) != 'nan' and int(announcedDate[0:2]) <= 4:
                # 组合成完成的公告年月日
                announcedDate = announcedYear + '-' + announcedDate
                print('%s %5.2f %s' % (code, eps, announcedDate), flush=True)

                # 如果当前股票不在需要返回的数据结构中，则添加到其中
                if code not in codesOfCachedReports:
                    codeReportDict[code] = []
                    codesOfCachedReports.add(code)
                # 将eps和公告日期添加到列表中
                codeReportDict[code].append({'eps': eps, 'announced_date': announcedDate})

    # 返回获取的数据
    return codeReportDict

"""
实现股票池选股逻辑，找到指定日期范围的候选股票
条件：0 < PE < 30, 按从小到大排序，剔除停牌后，取前100个；再平衡周期：7个交易日
:param begin_date: 开始日期
:param end_date: 结束日期
:return: tuple，再平衡的日期列表，以及一个dict(key: 再平衡日, value: 当期的股票列表)
"""
def stockPool(beginDate, endDate):
    # 获取财务数据
    codeReportDict = getCodeReports()

    # 股票池的再平衡周期
    rebalanceInterval = 7

    # 因为上证指数没有停牌不会缺数，所以用它作为交易日历
    szzsDf = ts.get_k_data('000001', index=True, start=beginDate, end=endDate)
    allDates = list(szzsDf['date'])
    print(szzsDf)

    # 调整日和其对应的股票
    rebalanceDateCodesDict = dict()
    rebalanceDates = []

    # 保存上一期的股票池
    lastPhaseCodes = []

    # 所有的交易日数
    datesCount = len(allDates)

    # 用再平衡周期作为步长循环
    for index in range(0, datesCount, rebalanceInterval):
        # 当前的调整日
        rebalanceDate = allDates[index]

        # 获取本期符合条件的备选股票
        thisPhaseOptionCodes = getOptionCodes(codeReportDict, rebalanceDate)

        # 本期入选的股票代码列表
        thisPhaseCodes = []

        # 找到在上一期的股票池，但是当前停牌的股票，保留在当期股票池中
        if len(lastPhaseCodes) > 0:
            for code in lastPhaseCodes:
                dailyK = ts.get_k_data(code, autype=None, start=rebalanceDate, end=rebalanceDate)
                if dailyK.size == 0:
                    thisPhaseCodes.append(code)
        print('上期停牌的股票：', flush=True)
        print(thisPhaseCodes, flush=True)

        # 剩余的位置用当前备选的股票
        optionSize = len(thisPhaseOptionCodes)
        if optionSize > (100 - len(thisPhaseCodes)):
            thisPhaseCodes += thisPhaseOptionCodes[0:100 - len(thisPhaseCodes)]
        else:
            thisPhaseCodes += thisPhaseOptionCodes

        # 当期股票池作为下次循环的上期股票池
        lastPhaseCodes = thisPhaseCodes

        # 保存到返回结果中
        rebalanceDateCodesDict[rebalanceDate] = thisPhaseCodes
        rebalanceDates.append(rebalanceDate)

        print('当前最终的备选票：%s' % rebalanceDate, flush=True)
        print(thisPhaseCodes, flush=True)

    return  rebalanceDates, rebalanceDateCodesDict

"""
找到某个调整日符合股票池条件的股票列表
:param code_report_dict: 股票对应的财报列表
:param rebalance_date: 再平衡日期
:return: 股票代码列表
"""
def getOptionCodes(codeReportDict, rebalanceDate):
    # 如果股票和每股收益的dict是空的，则重新获取
    reportCodes = list(codeReportDict.keys())
    if len(reportCodes) == 0:
        getCodeReports()
        reportCodes = list(codeReportDict.keys())

    # 找到当期符合条件的eps, 股票代码和eps
    codeEpsDict = dict()
    for code in reportCodes:
        # 因为财报是按照公告日期从早到晚排列的，所以顺序查找
        reports = codeReportDict[code]

        # 用来保存最后一个公告日期小于等于当前日期的财报
        lastReport = None

        for report in reports:
            announcedDate = report['announced_date']
            # 如果公告日期大于当前调整日期，则结束循环
            if announcedDate > rebalanceDate:
                break
            lastReport = report

        # 如果找到了正确时间范围的年报,并且eps大于0，才保留
        if lastReport is not None and lastReport['eps'] > 0:
            print('%s, %s, %s, %5.2f' %
                  (code, rebalanceDate, lastReport['announced_date'], lastReport['eps']), flush=True)

            codeEpsDict[code] = lastReport['eps']
    # 只在符合eps > 0的范围，计算PE，并筛选股票
    validatedCodes = list(codeEpsDict)

    dfPe = DataFrame(columns=['close', 'eps'])
    for code in validatedCodes:
        # 用不复权的价格
        dailyK = ts.get_k_data(code, autype=None, start=rebalanceDate, end=rebalanceDate)

        # 如果当前是停牌，就获取不到股份停牌，那么不参与排名
        if dailyK.size > 0:
            close = dailyK.loc[dailyK.index[0]]['close']
            dfPe.loc[code] = {'eps': codeEpsDict[code], 'close': close}
            print('%s %6.2f' % (code, close), flush=True)
        else:
            print('%s 停牌' % (code), flush=True)

    # 计算PE，重点表述为什么？ 重新讲解复权
    dfPe['pe'] = dfPe['close'] / dfPe['eps']

    # 从小到大排序
    dfPe.sort_values('pe', ascending=True, inplace=True)

    # 只保留小于30的数据
    dfPe = dfPe[dfPe['pe'] < 30]

    # 返回排名靠前的100只股票代码
    return list(dfPe.index)[0:100]

"""
统计股票池的收益
"""
def statisticStockPoolProfit():
    # 设定评测周期
    rebalanceDates, codesDict = stockPool('2018-01-01', '2018-09-12')

    # 用DataFrame保存收益
    dfProfit = DataFrame(columns=['profit', 'hs300'])

    dfProfit.loc[rebalanceDates[0]] = {'profit': 0, 'hs300': 0}

    # 获取沪深300在统计周期内的第一天的值
    hs300K = ts.get_k_data('000300', index=True, start=rebalanceDates[0], end=rebalanceDates[0])
    hs300BeginValue = hs300K.loc[hs300K.index[0]]['close']

    # 通过净值计算累计收益
    netValue = 1
    for index in range(1, len(rebalanceDates) - 1):
        lastRebalanceDate = rebalanceDates[index - 1]
        currentRebalanceDate = rebalanceDates[index]

        # 获取上一期的股票池
        codes = codesDict[lastRebalanceDate]

        # 统计当前的收益
        profitSum = 0

        # 参与统计收益的股票个数
        profitCodeCount = 0

        for code in codes:
            dailyKs = ts.get_k_data(code, autype='hfq', start=lastRebalanceDate, end=currentRebalanceDate)
            indexSize = dailyKs.index.size

            # 如果没有数据，则跳过， 长期停牌
            if indexSize == 0:
                continue

            # 买入价
            buyPrice = dailyKs.loc[dailyKs.index[0]]['close']

            # 卖出价
            sellPrice = dailyKs.loc[dailyKs.index[indexSize - 1]]['close']

            # 股票池内所有股票的收益
            profitSum += (sellPrice - buyPrice) / buyPrice
            profitCodeCount += 1

        profit = round(profitSum/profitCodeCount, 4)
        hs300KCurrent = ts.get_k_data('000300', index=True, start=currentRebalanceDate, end=currentRebalanceDate)
        hs300Close = hs300KCurrent.loc[hs300KCurrent.index[0]]['close']

        # 计算净值和累计收益
        netValue = netValue * (1 + profit)
        dfProfit.loc[currentRebalanceDate] = {
            'profit': round((netValue - 1) * 100, 4),
            'hs300': round((hs300Close - hs300BeginValue) * 100 / hs300BeginValue, 4)
        }
        print(dfProfit)

    # 绘制曲线
    dfProfit.plot(title='Stock Pool Profit Statistic', kind='line')

    # 显示图像
    plt.show()



if __name__ == '__main__':
    # getCodeReports()
    # stockPool('2008-01-01', '2018-09-12')
    statisticStockPoolProfit()